
#include "pch.h"

#include <SetupAPI.h>

HRESULT RetrieveDevicePath(
	_Out_bytecap_(BufLen) LPTSTR DevicePath,
	_In_				  ULONG  BufLen,
	_Out_opt_			  PBOOL  FailureDeviceNotFound
	);

/*++

Routine description:

Open all needed handles to interact with the device.

If the device has multiple USB interfaces, this function grants access to
only the first interface.

If multiple devices have the same device interface GUID, there is no
guarantee of which one will be returned.

Arguments:

DeviceData - Struct filled in by this function. The caller should use the
WinusbHandle to interact with the device, and must pass the struct to
CloseDevice when finished.

FailureDeviceNotFound - TRUE when failure is returned due to no devices
found with the correct device interface (device not connected, driver
not installed, or device is disabled in Device Manager); FALSE
otherwise.

Return value:

HRESULT

--*/
HRESULT OpenDevice(
	_Out_		PDEVICE_DATA DeviceData,
	_Out_opt_	PBOOL		 FailureDeviceNotFound
	)
{
	HRESULT hr = S_OK;
	BOOL bResult;

	DeviceData->HandlesOpen = FALSE;

	hr = RetrieveDevicePath(DeviceData->DevicePath,
		sizeof(DeviceData->DevicePath),
		FailureDeviceNotFound);

	if (FAILED(hr)) {
		return hr;
	}

	DeviceData->DeviceHandle = CreateFile(DeviceData->DevicePath,
		GENERIC_WRITE | GENERIC_READ,
		FILE_SHARE_WRITE | FILE_SHARE_READ,
		NULL,
		OPEN_EXISTING,
		FILE_ATTRIBUTE_NORMAL | FILE_FLAG_OVERLAPPED,
		NULL);

	if (INVALID_HANDLE_VALUE == DeviceData->DeviceHandle) {
		hr = HRESULT_FROM_WIN32(GetLastError());
		return hr;
	}

	bResult = WinUsb_Initialize(DeviceData->DeviceHandle,
		&DeviceData->WinusbHandle);

	if (FALSE == bResult) {
		hr = HRESULT_FROM_WIN32(GetLastError());
		CloseHandle(DeviceData->DeviceHandle);
	}

	DeviceData->HandlesOpen = TRUE;
	return hr;
}

/*++

Routine description:

Perform required cleanup when the device is no longer needed.

If OpenDevice failed, do nothing.

Arguments:

DeviceData - Struct filled in by OpenDevice

Return value:

None

--*/
VOID CloseDevice(
	_Inout_ PDEVICE_DATA DeviceData
	)
{
	if (FALSE == DeviceData->HandlesOpen) {
		return;
	}

	WinUsb_Free(DeviceData->WinusbHandle);
	CloseHandle(DeviceData->DeviceHandle);
	DeviceData->HandlesOpen = FALSE;
}

static const GUID USB_DEVICE_INTERFACE =
{ 0x88bae032, 0x5a81, 0x49f0, { 0xbc, 0x3d, 0xa4, 0xff, 0x13, 0x82, 0x16, 0xd6 } };

static const GUID USB_DEVICE =
{ 0xA5DCBF10, 0x6530, 0x11D2, { 0x90, 0x1F, 0x00, 0xC0, 0x4F, 0xB9, 0x51, 0xED } };

static const GUID MY_USB_DEVICE =
{ 0x408ff12f, 0x50f3, 0x42f9, { 0xad, 0x57, 0x54, 0x86, 0x48, 0xca, 0xe9, 0x55 } };

/*++

Routine description:

Retrieve the device path that can be used to open the WinUSB-based device.

If multiple devices have the same device interface GUID, there is no
guarantee of which one will be returned.

Arguments:

DevicePath - On successful return, the path of the device (use with CreateFile).

BufLen - The size of DevicePath's buffer, in bytes

FailureDeviceNotFound - TRUE when failure is returned due to no devices
found with the correct device interface (device not connected, driver
not installed, or device is disabled in Device Manager); FALSE
otherwise.

Return value:

HRESULT

--*/
HRESULT RetrieveDevicePath(
	_Out_bytecap_(BufLen) LPTSTR DevicePath,
	_In_				  ULONG  BufLen,
	_Out_opt_			  PBOOL  FailureDeviceNotFound
	)
{
	BOOL		bResult = FALSE;
	HDEVINFO	deviceInfo;
	SP_DEVICE_INTERFACE_DATA interfaceData;
	PSP_DEVICE_INTERFACE_DETAIL_DATA detailData = NULL;
	ULONG		length;
	ULONG		requiredLength = 0;
	HRESULT		hr;

	if (NULL != FailureDeviceNotFound) {
		*FailureDeviceNotFound = FALSE;
	}

	printf("RetrieveDevicePath \n");

	deviceInfo = SetupDiGetClassDevs(&USB_DEVICE,
		NULL,
		NULL,
		DIGCF_PRESENT | DIGCF_DEVICEINTERFACE);

	if (INVALID_HANDLE_VALUE == deviceInfo) {
		hr = HRESULT_FROM_WIN32(GetLastError());
		return hr;
	}

	interfaceData.cbSize = sizeof(SP_DEVICE_INTERFACE_DATA);

	bResult = SetupDiEnumDeviceInterfaces(deviceInfo,
		NULL,
		&USB_DEVICE,
		1,
		&interfaceData);

	if (FALSE == bResult) {
		if (ERROR_NO_MORE_ITEMS == GetLastError() &&
			NULL != FailureDeviceNotFound) {
			*FailureDeviceNotFound = TRUE;
		}

		hr = HRESULT_FROM_WIN32(GetLastError());
		SetupDiDestroyDeviceInfoList(deviceInfo);
		return hr;
	}

	bResult = SetupDiGetDeviceInterfaceDetail(deviceInfo,
		&interfaceData,
		NULL,
		0,
		&requiredLength,
		NULL);

	if (FALSE == bResult && ERROR_INSUFFICIENT_BUFFER != GetLastError()) {

		hr = HRESULT_FROM_WIN32(GetLastError());
		SetupDiDestroyDeviceInfoList(deviceInfo);
		return hr;
	}

	detailData = (PSP_DEVICE_INTERFACE_DETAIL_DATA)LocalAlloc(LMEM_FIXED,
		requiredLength);

	if (NULL == detailData) {

		hr = E_OUTOFMEMORY;
		SetupDiDestroyDeviceInfoList(deviceInfo);
		return hr;
	}

	detailData->cbSize = sizeof(SP_DEVICE_INTERFACE_DETAIL_DATA);
	length = requiredLength;

	bResult = SetupDiGetDeviceInterfaceDetail(deviceInfo,
		&interfaceData,
		detailData,
		length,
		&requiredLength,
		NULL);

	if (FALSE == bResult) {
		
		hr = HRESULT_FROM_WIN32(GetLastError());
		LocalFree(detailData);
		SetupDiDestroyDeviceInfoList(deviceInfo);
		return hr;
	}

	hr = StringCbCopy(DevicePath,
		BufLen,
		detailData->DevicePath);
	LocalFree(detailData);
	SetupDiDestroyDeviceInfoList(deviceInfo);

	return hr;
}

BOOL GetUSBDeviceSpeed(
	WINUSB_INTERFACE_HANDLE hDeviceHandle,
	UCHAR* pDeviceSpeed
	)
{
	if (!pDeviceSpeed || INVALID_HANDLE_VALUE == hDeviceHandle)
	{
		return FALSE;
	}

	BOOL bResult = TRUE;
	ULONG length = sizeof(UCHAR);

	bResult = WinUsb_QueryDeviceInformation(hDeviceHandle,
		DEVICE_SPEED,
		&length,
		pDeviceSpeed);
	if (!bResult)
	{
		printf("Error getting device speed: %d.\n", GetLastError());
		return bResult;
	}

	if (*pDeviceSpeed == LowSpeed)
	{
		printf("Device speed: %d (Low speed).\n", *pDeviceSpeed);
	}
	else if (*pDeviceSpeed == FullSpeed)
	{
		printf("Device speed: %d (Full speed).\n", *pDeviceSpeed);
	}
	else if (*pDeviceSpeed == HighSpeed)
	{
		printf("Device speed: %d (High speed).\n", *pDeviceSpeed);
	}
	
	return bResult;
}

BOOL QueryDeviceEndpoints(
	WINUSB_INTERFACE_HANDLE hDeviceHandle,
	PIPE_ID* pipeid
	)
{
	if (INVALID_HANDLE_VALUE == hDeviceHandle)
	{
		return FALSE;
	}

	BOOL bResult = TRUE;

	USB_INTERFACE_DESCRIPTOR InterfaceDescriptor;
	ZeroMemory(&InterfaceDescriptor, sizeof(USB_INTERFACE_DESCRIPTOR));

	WINUSB_PIPE_INFORMATION Pipe;
	ZeroMemory(&Pipe, sizeof(WINUSB_PIPE_INFORMATION));

	bResult = WinUsb_QueryInterfaceSettings(hDeviceHandle, 0, &InterfaceDescriptor);

	if (bResult)
	{
		for (int index = 0; index < InterfaceDescriptor.bNumEndpoints; index++)
		{
			bResult = WinUsb_QueryPipe(hDeviceHandle, 0, index, &Pipe);

			if (bResult)
			{
				if (UsbdPipeTypeControl == Pipe.PipeType)
				{
					printf("Endpoint index: %d Pipe type: %d Control Pipe ID: %d.\n", index, Pipe.PipeType, Pipe.PipeId);
				}
				else if (UsbdPipeTypeIsochronous == Pipe.PipeType)
				{
					printf("Endpoint index: %d Pipe type: %d IsoChronous Pipe ID: %d.\n", index, Pipe.PipeType, Pipe.PipeId);
				}
				else if (UsbdPipeTypeBulk == Pipe.PipeType)
				{
					if (USB_ENDPOINT_DIRECTION_IN(Pipe.PipeId))
					{
						printf("Endpoint index: %d Pipe Type: %d Bulk Pipe ID: %c.\n", index, Pipe.PipeType, Pipe.PipeId);
						pipeid->PipeInId = Pipe.PipeId;
					}
					else if (USB_ENDPOINT_DIRECTION_OUT(Pipe.PipeId))
					{
						printf("Endpoint index: %d Pipe Type: %d Bulk Pipe ID: %c.\n", index, Pipe.PipeType, Pipe.PipeId);
					}
				}
				else if (UsbdPipeTypeInterrupt == Pipe.PipeType)
				{
					printf("Endpoint index: %d Pipe Type: %d Interrupt Pipe ID: %d.\n", index, Pipe.PipeType, Pipe.PipeId);
				}
			}
			else
			{
				continue;
			}
		}
	}

	return bResult;
}


